from pymodbus.client import AsyncModbusSerialClient, AsyncModbusTcpClient
from pymodbus.client import ModbusSerialClient, ModbusTcpClient
from pymodbus.exceptions import ModbusException
from src.utils.loguru_func import modbus_logger
import time
import json


class ModbusToolAsync:
    def __init__(self, client_type='rtu', **kwargs):
        """
        初始化 Modbus 客户端
        :param client_type: 客户端类型，'rtu' 或 'tcp'
        :param kwargs: 客户端连接参数
        """
        if client_type == 'rtu':
            self.client = AsyncModbusSerialClient(**kwargs)
        elif client_type == 'tcp':
            self.client = AsyncModbusTcpClient(**kwargs)
        else:
            raise ValueError("client_type 必须为 'rtu' 或 'tcp'")

    async def connect(self):
        """连接到 Modbus 从站"""
        try:
            await self.client.connect()
            print("成功连接到 Modbus 从站")
            return True
        except ModbusException as e:
            print(f"连接失败: {e}")
            return False

    async def read_holding_registers(self, address, count, slave=1):
        """
        读取保持寄存器
        :param address: 起始地址
        :param count: 读取数量
        :param slave: 从站地址
        :return: 读取结果或 None
        """
        try:
            result = await self.client.read_holding_registers(address=address, count=count, slave=slave)
            if not result.isError():
                return result.registers
            else:
                print(f"读取保持寄存器错误: {result}")
        except ModbusException as e:
            print(f"读取保持寄存器时发生异常: {e}")
        return None

    async def write_single_register(self, address, value, slave=1):
        """
        写入单个保持寄存器
        :param address: 寄存器地址
        :param value: 写入的值
        :param slave: 从站地址
        :return: 写入成功返回 True，失败返回 False
        """
        try:
            result = await self.client.write_register(address=address, value=value, slave=slave)
            if not result.isError():
                print("写入单个保持寄存器成功")
                return True
            else:
                print(f"写入单个保持寄存器错误: {result}")
        except ModbusException as e:
            print(f"写入单个保持寄存器时发生异常: {e}")
        return False

    async def write_multiple_registers(self, address, values, slave=1):
        """
        写入多个保持寄存器
        :param address: 起始地址
        :param values: 写入的值列表
        :param slave: 从站地址
        :return: 写入成功返回 True，失败返回 False
        """
        try:
            result = await self.client.write_registers(address=address, values=values, slave=slave)
            if not result.isError():
                print("写入多个保持寄存器成功")
                return True
            else:
                print(f"写入多个保持寄存器错误: {result}")
        except ModbusException as e:
            print(f"写入多个保持寄存器时发生异常: {e}")
        return False

    async def close(self):
        """关闭 Modbus 连接"""
        self.client.close()
        print("Modbus 连接已关闭")


class ModbusToolSync:
    def __init__(self, client_type='rtu', **kwargs):
        """
        初始化 Modbus 客户端
        :param client_type: 客户端类型，'rtu' 或 'tcp'
        :param kwargs: 客户端连接参数
        """
        if client_type == 'rtu':
            self.client = ModbusSerialClient(**kwargs)
        elif client_type == 'tcp':
            self.client = ModbusTcpClient(**kwargs)
        else:
            raise ValueError("client_type 必须为 'rtu' 或 'tcp'")

    def connect(self):
        """连接到 Modbus 从站"""
        try:
            if self.client.connect():
                modbus_logger.debug("成功连接到 Modbus 从站")
                return True
            else:
                modbus_logger.debug("连接失败")
                return False
        except ModbusException as e:
            modbus_logger.debug(f"连接失败: {e}")
            return False

    def read_coils(self, address, count, slave=1):
        """
        01H 读线圈状态
        :param address: 起始地址
        :param count: 读取数量
        :param slave: 从站地址
        :return: 读取结果或 None
        """
        try:
            result = self.client.read_coils(
                address=address, count=count, slave=slave)
            if not result.isError():
                modbus_logger.debug(f"读线圈状态成功{result.bits}")
                return result.bits
            else:
                modbus_logger.debug(f"读线圈状态错误: {result}")
        except ModbusException as e:
            modbus_logger.debug(f"读线圈状态时发生异常: {e}")
        return None

    def read_discrete_inputs(self, address, count, slave=1):
        """
        02H 读离散输入状态
        :param address: 起始地址
        :param count: 读取数量
        :param slave: 从站地址
        :return: 读取结果或 None
        """
        try:
            result = self.client.read_discrete_inputs(
                address=address, count=count, slave=slave)
            if not result.isError():
                modbus_logger.debug(f"读离散输入状态成功{result.bits}")
                return result.bits
            else:
                modbus_logger.debug(f"读离散输入状态错误: {result}")
        except ModbusException as e:
            modbus_logger.debug(f"读离散输入状态时发生异常: {e}")
        return None

    def read_holding_registers(self, address, count, slave=1):
        """
        03H 读保持寄存器
        :param address: 起始地址
        :param count: 读取数量
        :param slave: 从站地址
        :return: 读取结果或 None
        """
        try:
            result = self.client.read_holding_registers(
                address=address, count=count, slave=slave)
            if not result.isError():
                modbus_logger.debug(f"读取保持寄存器成功{result.registers}")
                return result.registers
            else:
                modbus_logger.debug(f"读取保持寄存器错误: {result}")
        except ModbusException as e:
            modbus_logger.debug(f"读取保持寄存器时发生异常: {e}")
        return None

    def read_input_registers(self, address, count, slave=1):
        """
        04H 读输入寄存器
        :param address: 起始地址
        :param count: 读取数量
        :param slave: 从站地址
        :return: 读取结果或 None
        """
        try:
            result = self.client.read_input_registers(
                address=address, count=count, slave=slave)
            if not result.isError():
                modbus_logger.debug(f"读输入寄存器成功{result.registers}")
                return result.registers
            else:
                modbus_logger.debug(f"读输入寄存器错误: {result}")
        except ModbusException as e:
            modbus_logger.debug(f"读输入寄存器时发生异常: {e}")
        return None

    def write_single_coil(self, address, value, slave=1):
        """
        05H 写单个线圈
        :param address: 线圈地址
        :param value: 写入的值，True 或 False
        :param slave: 从站地址
        :return: 写入成功返回 True，失败返回 False
        """
        try:
            result = self.client.write_coil(
                address=address, value=value, slave=slave)
            if not result.isError():
                modbus_logger.debug(f"写单个线圈成功{result}")
                return True
            else:
                modbus_logger.debug(f"写入单个线圈错误: {result}")
        except ModbusException as e:
            modbus_logger.debug(f"写入单个线圈时发生异常: {e}")
        return False

    def write_single_register(self, address, value, slave=1):
        """
        06H 写单个保持寄存器
        :param address: 寄存器地址
        :param value: 写入的值
        :param slave: 从站地址
        :return: 写入成功返回 True，失败返回 False
        """
        try:
            result = self.client.write_register(
                address=address, value=value, slave=slave)
            if not result.isError():
                modbus_logger.debug(f"写单个保持寄存器成功{result}")
                return True
            else:
                modbus_logger.debug(f"写入单个保持寄存器错误: {result}")
        except ModbusException as e:
            modbus_logger.debug(f"写入单个保持寄存器时发生异常: {e}")
        return False

    def write_multiple_coils(self, address, values, slave=1):
        """
        0FH 写多个线圈
        :param address: 起始地址
        :param values: 写入的值列表，元素为 True 或 False
        :param slave: 从站地址
        :return: 写入成功返回 True，失败返回 False
        """
        try:
            result = self.client.write_coils(
                address=address, values=values, slave=slave)
            if not result.isError():
                modbus_logger.debug(f"写多个线圈成功{result}")
                return True
            else:
                modbus_logger.debug(f"写入多个线圈错误: {result}")
        except ModbusException as e:
            modbus_logger.debug(f"写入多个线圈时发生异常: {e}")
        return False

    def write_multiple_registers(self, address, values, slave=1):
        """
        10H 写多个保持寄存器
        :param address: 起始地址
        :param values: 写入的值列表
        :param slave: 从站地址
        :return: 写入成功返回 True，失败返回 False
        """
        try:
            result = self.client.write_registers(
                address=address, values=values, slave=slave)
            if not result.isError():
                modbus_logger.debug(f"写多个保持寄存器成功{result}")
                return True
            else:
                modbus_logger.debug(f"写入多个保持寄存器错误: {result}")
        except ModbusException as e:
            modbus_logger.debug(f"写入多个保持寄存器时发生异常: {e}")
        return False

    def disConnect(self):
        """关闭 Modbus 连接"""
        self.client.close()
        modbus_logger.debug("Modbus 连接已关闭")


class Huikong_Output:
    # [x] 汇控IO板的操作封装，待完善
    def __init__(self, port: str, baudrate: int = 9600, bytesize: int = 8, parity: str = 'N', stopbits: int = 1, timeout: float = 1):
        self.modbus_rtu = ModbusToolSync(
            port=port,
            baudrate=baudrate,
            bytesize=bytesize,
            parity=parity,
            stopbits=stopbits,
            timeout=timeout
        )

        # FIXME 汇控IO板的从站号，写死后不去修改
        self.slave = 0x01
        self.set_trigger_time()

        self.connect()

    def set_trigger_time(self):
        with open('config/trigger_time.json', 'r', encoding='utf-8') as f:
            trigger_time_dict = json.load(f)

        # 配置顺序开关通道的时间
        self.trigger_time = trigger_time_dict['trigger_time']
        self.interval_time = trigger_time_dict['interval_time']
        self.trigger_channels = trigger_time_dict['trigger_channels']
        modbus_logger.debug(self.trigger_channels)

    def connect(self):
        self.modbus_rtu.connect()

    def disConnect(self):
        self.modbus_rtu.disConnect()

    def trigger_on_channel(self, output_channel: int) -> bool:

        # 通道号，从1开始计数
        # 写单个寄存器，址为 0x0000 + 通道号 - 1
        return self.modbus_rtu.write_single_register(
            address=0x0000 + output_channel - 1, value=1, slave=self.slave)

    def trigger_off_channel(self, output_channel: int):
        # 通道号，从1开始计数
        # 写单个寄存器，址为 0x0000 + 通道号 - 1
        return self.modbus_rtu.write_single_register(
            address=0x0000 + output_channel - 1, value=0, slave=self.slave)

    # 按顺序触发1-20号通道，每个触发500ms，间隔100ms，再触发下一个
    def trigger_channels_in_order(self):
        trigger_counts_result = {}
        trigger_counts_result['on'] = 0
        trigger_counts_result['off'] = 0

        for i in self.trigger_channels:
            if self.trigger_on_channel(i):
                trigger_counts_result['on'] += 1
            time.sleep(self.trigger_time)
            if self.trigger_off_channel(i):
                trigger_counts_result['off'] += 1
            time.sleep(self.interval_time)

        return trigger_counts_result


async def test_async():
    # RTU 示例
    rtu_tool = ModbusToolAsync(
        client_type='rtu',
        port='COM3',
        baudrate=9600,
        bytesize=8,
        parity='N',
        stopbits=1,
        timeout=1
    )
    await rtu_tool.connect()
    # 读取保持寄存器
    result = await rtu_tool.read_holding_registers(address=0, count=10)
    if result:
        print("读取保持寄存器结果:", result)
    # 写入单个寄存器
    await rtu_tool.write_single_register(address=1, value=111)
    # 写入多个寄存器
    await rtu_tool.write_multiple_registers(address=1, values=[312, 313, 314])
    # 再次读取验证
    result = await rtu_tool.read_holding_registers(address=0, count=10)
    if result:
        print("再次读取保持寄存器结果:", result)
    await rtu_tool.close()

    # TCP 示例
    # tcp_tool = ModbusTool(
    #     client_type='tcp',
    #     host='127.0.0.1',
    #     port=502
    # )
    # await tcp_tool.connect()
    # # 这里可以添加 TCP 相关的读写操作
    # await tcp_tool.close()

# asyncio.run(test_async())


def test_sync():
    # RTU 示例
    rtu_tool = ModbusToolSync(
        client_type='rtu',
        port='COM3',
        baudrate=9600,
        bytesize=8,
        parity='N',
        stopbits=1,
        timeout=1
    )
    if rtu_tool.connect():
        # 读取保持寄存器
        result = rtu_tool.read_holding_registers(address=1, count=10)
        if result:
            print("读取保持寄存器结果:", result)
        # 写入单个寄存器
        rtu_tool.write_single_register(address=1, value=111)
        # 写入多个寄存器
        rtu_tool.write_multiple_registers(
            address=1, values=[312, 313, 314])
        # 再次读取验证
        result = rtu_tool.read_holding_registers(address=0, count=10)
        if result:
            print("再次读取保持寄存器结果:", result)
        rtu_tool.disConnect()
# test_sync()
